{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5c95e9bf",
   "metadata": {},
   "source": [
    "# Plan-and-Execute\n",
    "\n",
    "이 튜토리얼에서는 \"plan-and-execute\" 스타일의 에이전트를 만드는 방법을 소개하고, 이를 [LangGraph](https://langchain-ai.github.io/langgraph/)를 활용하여 구현하는 과정을 단계별로 설명합니다.  \n",
    "\"plan-and-execute\" 전략은 복잡한 작업을 수행할 때 장기적인 계획을 먼저 수립한 뒤, 해당 계획을 단계별로 실행하며 필요에 따라 다시 계획을 재수정하는 접근법입니다.\n",
    "\n",
    "![](./assets/langgraph-plan-and-execute.png)\n",
    "\n",
    "---\n",
    "\n",
    "## Plan-and-Execute란 무엇인가?\n",
    "\n",
    "\"plan-and-execute\"는 다음과 같은 특징을 갖는 접근 방식입니다.\n",
    "\n",
    "- **장기 계획 수립**: 복잡한 작업을 수행하기 전에 큰 그림을 그리는 장기 계획을 수립합니다.\n",
    "- **단계별 실행 및 재계획**: 세운 계획을 단계별로 실행하고, 각 단계가 완료될 때마다 계획이 여전히 유효한지 검토한 뒤 수정할 수 있습니다.\n",
    "  \n",
    "이 방식은 [Plan-and-Solve 논문](https://arxiv.org/abs/2305.04091)과 [Baby-AGI 프로젝트](https://github.com/yoheinakajima/babyagi)에서 영감을 받았습니다. 전통적인 [ReAct 스타일](https://arxiv.org/abs/2210.03629)의 에이전트는 한 번에 한 단계씩 생각하는 반면, \"plan-and-execute\"는 명시적이고 장기적인 계획을 강조합니다.\n",
    "\n",
    "**장점**:\n",
    "1. **명시적인 장기 계획**: 강력한 LLM조차도 한 번에 장기 계획을 처리하는 데 어려움을 겪을 수 있습니다. 명시적으로 장기 계획을 수립함으로써, 보다 안정적인 진행이 가능합니다.\n",
    "2. **효율적인 모델 사용**: 계획 단계에서는 더 큰/강력한 모델을 사용하고, 실행 단계에서는 상대적으로 작은/약한 모델을 사용함으로써 자원 소비를 최적화할 수 있습니다.\n",
    "\n",
    "---\n",
    "\n",
    "**주요 내용**\n",
    "\n",
    "- **도구 정의**: 사용할 도구 정의\n",
    "- **실행 에이전트 정의**: 실제 작업을 실행하는 에이전트 생성\n",
    "- **상태 정의**: 에이전트의 상태 정의\n",
    "- **계획 단계**: 장기 계획을 세우는 단계 생성\n",
    "- **재계획 단계**: 작업 진행 상황에 따라 계획을 재수정하는 단계 생성\n",
    "- **그래프 생성 및 실행**: 이러한 단계들을 연결하는 그래프 생성 및 실행\n",
    "\n",
    "---\n",
    "\n",
    "**참고**\n",
    "\n",
    "- [LangGraph 공식 문서](https://langchain-ai.github.io/langgraph/)  \n",
    "- [Plan-and-Solve 논문](https://arxiv.org/abs/2305.04091)  \n",
    "- [Baby-AGI 프로젝트](https://github.com/yoheinakajima/babyagi)  \n",
    "- [ReAct 논문](https://arxiv.org/abs/2210.03629)\n",
    "\n",
    "지금부터는 각 단계를 따라가며 \"plan-and-execute\" 에이전트를 LangGraph로 구현하는 방법을 자세히 알아보겠습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e99fefd6",
   "metadata": {},
   "source": [
    "## 환경 설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb1ad59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b4bac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Use-Cases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6afbb11",
   "metadata": {},
   "source": [
    "실습에 활용할 모델명 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afe4377f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.models import get_model_name, LLMs\n",
    "\n",
    "# 모델명 정의\n",
    "MODEL_NAME = get_model_name(LLMs.GPT4o)\n",
    "print(MODEL_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0856a092",
   "metadata": {},
   "source": [
    "## 도구 정의\n",
    "\n",
    "사용할 도구를 먼저 정의합니다. 이 간단한 예제에서는 `Tavily`를 통해 제공되는 내장 검색 도구를 사용할 것입니다. 그러나 직접 도구를 만드는 것도 매우 쉽습니다. \n",
    "\n",
    "자세한 내용은 [도구(Tools)](https://wikidocs.net/262582) 문서를 참조하십시오."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4dc76788",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.tools import TavilySearch\n",
    "\n",
    "# Tavily 검색 도구 초기화\n",
    "tools = [TavilySearch(max_results=3)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d230b5c",
   "metadata": {},
   "source": [
    "## 작업 실행 에이전트 정의\n",
    "\n",
    "이제 작업을 실행할 `execution agent`를 생성합니다. \n",
    "\n",
    "이 예제에서는 각 작업에 동일한 `execution agent`를 사용할 예정이지만, 반드시 이렇게 해야 하는 것은 아닙니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f70cf9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "# 프롬프트 정의\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant. Answer in Korean.\",\n",
    "        ),\n",
    "        (\"human\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "# LLM 정의\n",
    "llm = ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "\n",
    "# ReAct 에이전트 생성\n",
    "agent_executor = create_react_agent(llm, tools, state_modifier=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab1307f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 에이전트 실행\n",
    "agent_executor.invoke(\n",
    "    {\"messages\": [(\"user\", \"랭체인 한국어 튜토리얼에 대해서 설명해줘\")]}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05f6efaa",
   "metadata": {},
   "source": [
    "## 상태 정의\n",
    "\n",
    "- `input`: 사용자의 입력\n",
    "- `plan`: 현재 계획\n",
    "- `past_steps`: 이전에 실행한 계획과 실행 결과\n",
    "- `response`: 최종 응답"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "27c44393",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, List, Tuple\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "# 상태 정의\n",
    "class PlanExecute(TypedDict):\n",
    "    input: Annotated[str, \"User's input\"]\n",
    "    plan: Annotated[List[str], \"Current plan\"]\n",
    "    past_steps: Annotated[List[Tuple], operator.add]\n",
    "    response: Annotated[str, \"Final response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c5d026",
   "metadata": {},
   "source": [
    "## 계획(Plan) 단계\n",
    "\n",
    "이제 **계획 단계**를 생성하는 방법을 고려해 보겠습니다. 이 단계에서는 `function calling` 을 사용하여 계획을 수립합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ad37fbe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from typing import List\n",
    "\n",
    "\n",
    "# Plan 모델 정의\n",
    "class Plan(BaseModel):\n",
    "    \"\"\"Sorted steps to execute the plan\"\"\"\n",
    "\n",
    "    steps: Annotated[List[str], \"Different steps to follow, should be in sorted order\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "726052a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# 계획 수립을 위한 프롬프트 템플릿 생성\n",
    "planner_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"\"\"For the given objective, come up with a simple step by step plan. \\\n",
    "This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \\\n",
    "The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.\n",
    "Answer in Korean.\"\"\",\n",
    "        ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "planner = planner_prompt | ChatOpenAI(\n",
    "    model=MODEL_NAME, temperature=0\n",
    ").with_structured_output(Plan)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0fd64cd",
   "metadata": {},
   "source": [
    "`planner` 를 실행하여 계획을 수립한 결과를 확인합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c782515c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Planner 실행\n",
    "planner.invoke(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            (\n",
    "                \"user\",\n",
    "                \"LangGraph 의 핵심 장단점과 LangGraph 를 사용하는 이유는 무엇인가?\",\n",
    "            )\n",
    "        ]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef28ddf5",
   "metadata": {},
   "source": [
    "## 재계획(Re-Plan) 단계\n",
    "\n",
    "이제 이전 단계의 결과를 바탕으로 계획을 다시 수립하는 단계를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c80e74e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "\n",
    "\n",
    "class Response(BaseModel):\n",
    "    \"\"\"Response to user.\"\"\"\n",
    "\n",
    "    # 사용자 응답\n",
    "    response: str\n",
    "\n",
    "\n",
    "class Act(BaseModel):\n",
    "    \"\"\"Action to perform.\"\"\"\n",
    "\n",
    "    # 수행할 작업: \"Response\", \"Plan\". 사용자에게 응답할 경우 Response 사용, 추가 도구 사용이 필요할 경우 Plan 사용\n",
    "    action: Union[Response, Plan] = Field(\n",
    "        description=\"Action to perform. If you want to respond to user, use Response. \"\n",
    "        \"If you need to further use tools to get the answer, use Plan.\"\n",
    "    )\n",
    "\n",
    "\n",
    "# 계획을 재수립하기 위한 프롬프트 정의\n",
    "replanner_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"For the given objective, come up with a simple step by step plan. \\\n",
    "This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \\\n",
    "The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.\n",
    "\n",
    "Your objective was this:\n",
    "{input}\n",
    "\n",
    "Your original plan was this:\n",
    "{plan}\n",
    "\n",
    "You have currently done the follow steps:\n",
    "{past_steps}\n",
    "\n",
    "Update your plan accordingly. If no more steps are needed and you can return to the user, then respond with that. Otherwise, fill out the plan. Only add steps to the plan that still NEED to be done. Do not return previously done steps as part of the plan.\n",
    "\n",
    "Answer in Korean.\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "# Replanner 생성\n",
    "replanner = replanner_prompt | ChatOpenAI(\n",
    "    model=MODEL_NAME, temperature=0\n",
    ").with_structured_output(Act)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88a7250",
   "metadata": {},
   "source": [
    "## 그래프 생성\n",
    "\n",
    "이제 그래프를 생성할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "df12d945",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "\n",
    "# 사용자 입력을 기반으로 계획을 생성하고 반환\n",
    "def plan_step(state: PlanExecute):\n",
    "    plan = planner.invoke({\"messages\": [(\"user\", state[\"input\"])]})\n",
    "    # 생성된 계획의 단계 리스트 반환\n",
    "    return {\"plan\": plan.steps}\n",
    "\n",
    "\n",
    "# 에이전트 실행기를 사용하여 주어진 작업을 수행하고 결과를 반환\n",
    "def execute_step(state: PlanExecute):\n",
    "    plan = state[\"plan\"]\n",
    "    # 계획을 문자열로 변환하여 각 단계에 번호를 매김\n",
    "    plan_str = \"\\n\".join(f\"{i+1}. {step}\" for i, step in enumerate(plan))\n",
    "    task = plan[0]\n",
    "    # 현재 실행할 작업을 포맷팅하여 에이전트에 전달\n",
    "    task_formatted = f\"\"\"For the following plan:\n",
    "{plan_str}\\n\\nYou are tasked with executing [step 1. {task}].\"\"\"\n",
    "    # 에이전트 실행기를 통해 작업 수행 및 결과 수신\n",
    "    agent_response = agent_executor.invoke({\"messages\": [(\"user\", task_formatted)]})\n",
    "    # 이전 단계와 그 결과를 포함하는 딕셔너리 반환\n",
    "    return {\n",
    "        \"past_steps\": [(task, agent_response[\"messages\"][-1].content)],\n",
    "    }\n",
    "\n",
    "\n",
    "# 이전 단계의 결과를 바탕으로 계획을 업데이트하거나 최종 응답을 반환\n",
    "def replan_step(state: PlanExecute):\n",
    "    output = replanner.invoke(state)\n",
    "    # 응답이 사용자에게 반환될 경우\n",
    "    if isinstance(output.action, Response):\n",
    "        return {\"response\": output.action.response}\n",
    "    # 추가 단계가 필요할 경우 계획의 단계 리스트 반환\n",
    "    else:\n",
    "        next_plan = output.action.steps\n",
    "        if len(next_plan) == 0:\n",
    "            return {\"response\": \"No more steps needed.\"}\n",
    "        else:\n",
    "            return {\"plan\": next_plan}\n",
    "\n",
    "\n",
    "# 에이전트의 실행 종료 여부를 결정하는 함수\n",
    "def should_end(state: PlanExecute):\n",
    "    if \"response\" in state and state[\"response\"]:\n",
    "        return \"final_report\"\n",
    "    else:\n",
    "        return \"execute\"\n",
    "\n",
    "\n",
    "final_report_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"You are given the objective and the previously done steps. Your task is to generate a final report in markdown format.\n",
    "Final report should be written in professional tone.\n",
    "\n",
    "Your objective was this:\n",
    "\n",
    "{input}\n",
    "\n",
    "Your previously done steps(question and answer pairs):\n",
    "\n",
    "{past_steps}\n",
    "\n",
    "Generate a final report in markdown format. Write your response in Korean.\"\"\"\n",
    ")\n",
    "\n",
    "final_report = (\n",
    "    final_report_prompt\n",
    "    | ChatOpenAI(model=MODEL_NAME, temperature=0)\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "\n",
    "def generate_final_report(state: PlanExecute):\n",
    "    past_steps = \"\\n\\n\".join(\n",
    "        [\n",
    "            f\"Question: {past_step[0]}\\n\\nAnswer: {past_step[1]}\\n\\n####\"\n",
    "            for past_step in state[\"past_steps\"]\n",
    "        ]\n",
    "    )\n",
    "    response = final_report.invoke({\"input\": state[\"input\"], \"past_steps\": past_steps})\n",
    "    return {\"response\": response}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff7ebe60",
   "metadata": {},
   "source": [
    "## 그래프 생성\n",
    "\n",
    "이제 지금까지 정의한 노드를 연결하여 그래프를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "16abdcf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "\n",
    "# 그래프 생성\n",
    "workflow = StateGraph(PlanExecute)\n",
    "\n",
    "# 노드 정의\n",
    "workflow.add_node(\"planner\", plan_step)\n",
    "workflow.add_node(\"execute\", execute_step)\n",
    "workflow.add_node(\"replan\", replan_step)\n",
    "workflow.add_node(\"final_report\", generate_final_report)\n",
    "\n",
    "# 엣지 정의\n",
    "workflow.add_edge(START, \"planner\")\n",
    "workflow.add_edge(\"planner\", \"execute\")\n",
    "workflow.add_edge(\"execute\", \"replan\")\n",
    "workflow.add_edge(\"final_report\", END)\n",
    "\n",
    "# 조건부 엣지: replan 후 종료 여부를 결정하는 함수 사용\n",
    "workflow.add_conditional_edges(\n",
    "    \"replan\",\n",
    "    should_end,\n",
    "    {\"execute\": \"execute\", \"final_report\": \"final_report\"},\n",
    ")\n",
    "\n",
    "# 그래프 컴파일\n",
    "app = workflow.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd850b8c",
   "metadata": {},
   "source": [
    "그래프를 시각화 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee7d379",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "visualize_graph(app, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d266cb",
   "metadata": {},
   "source": [
    "## 그래프 실행"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6323963b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.messages import invoke_graph, random_uuid\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "config = RunnableConfig(recursion_limit=50, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "inputs = {\n",
    "    \"input\": \"Modular RAG 가 기존의 Naive RAG 와 어떤 차이가 있는지와 production level 에서 사용하는 이점을 설명해줘\"\n",
    "}\n",
    "\n",
    "invoke_graph(app, inputs, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e139d290",
   "metadata": {},
   "outputs": [],
   "source": [
    "snapshot = app.get_state(config).values\n",
    "print(snapshot[\"response\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611af101",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "Markdown(snapshot[\"response\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fde7887",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(snapshot[\"response\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
